#!/usr/bin/env python
# encoding: utf-8
"""
@author: youfeng
@email: youfeng243@163.com
@license: Apache Licence 
@file: queue_mq_thread.py
@time: 2016/12/8 13:45
"""
import json
import threading
import time
from Queue import Queue

from beanstalkc import SocketError
from thrift.protocol.TBinaryProtocol import TBinaryProtocol
from thrift.transport.TTransport import TMemoryBuffer

from common.pybeanstalk import PyBeanstalk
from common.tools import build_hzpost_url
from config.conf import beanstalk_server_conf
from thriftRPC.ttypes import DownLoadRsp


class MqQueueThread(threading.Thread):
    def __init__(self, server_conf=None, log=None, is_open=True):
        threading.Thread.__init__(self)
        self.daemon = True

        self.log = log

        self.is_open = is_open

        # 输送队列
        self.queue = Queue()

        # 获取消息队列配置
        self.server_conf = server_conf if server_conf is not None else beanstalk_server_conf

        # 消息队列
        self.beanstalk = PyBeanstalk(self.server_conf['host'], self.server_conf['port'])
        self.output_tube = self.server_conf['tube']

    def __del__(self):
        self.log.info('消息队列线程退出...')

    def push(self, msg, post_data=None):

        url = msg.get('url', '')
        content = msg.get('text', '')
        if url == '' or content == '':
            return

        self.push_msg(url, content, post_data=post_data)

    def push_msg(self, url, text, post_data=None, company=None, year=None):
        # 如果是post请求 则需要进行url加密
        if isinstance(post_data, dict):
            url = build_hzpost_url(url, post_data)

        obj = self.__download_rsp(url, content=text.encode("UTF-8"), company=company, year=year)
        obj_str = self.__to_string(obj)
        if obj_str is None:
            return

        self.queue.put_nowait(obj_str)

    def close(self):
        self.queue.put_nowait('@@##$$')
        self.log.info('发送线程退出指令...')

    @staticmethod
    def __download_rsp(url, content, company=None, year=None, content_type="html/text"):
        obj = DownLoadRsp()
        obj.url = url
        obj.elapsed = 50
        obj.content = content
        obj.page_size = len(obj.content)
        obj.content_type = content_type
        obj.status = 0
        obj.http_code = 200
        obj.download_time = int(time.time())

        if year is not None:
            site_recode_id = company + '_' + year
        else:
            site_recode_id = company

        # 如果没有company字段则不发送这个字段
        if company is not None:
            obj.data_extends = json.dumps({"company": company, '_site_record_id': site_recode_id})
        return obj

    def __to_string(self, obj):
        str_parse = None
        try:
            memory_b = TMemoryBuffer()
            t_binary_protocol_b = TBinaryProtocol(memory_b)
            obj.write(t_binary_protocol_b)
            str_parse = memory_b.getvalue()
        except EOFError as e:
            self.log.exception(e)
        return str_parse

    def run(self):
        self.log.info('开始运行消息队列...')
        self.log.info('当前消息队列大小: size = {size}'.format(size=self.queue.qsize()))
        while True:
            try:
                msg = self.queue.get()

                if msg == '@@##$$':
                    break

                while True:
                    try:
                        # 志强解析模块未上线, 还不能发送数据到消息队列
                        if self.is_open:
                            self.beanstalk.put(self.output_tube, msg)
                        # 测试控制速度
                        # time.sleep(5)
                        break
                    except SocketError as e:
                        time.sleep(10)
                        self.beanstalk.reconnect()
                        self.log.warn("reconnect beanstalk...")
                        self.log.exception(e)
            except Exception as e:
                self.log.info('当前队列大小: size = {size}'.format(size=self.queue.qsize()))
                self.log.exception(e)
                time.sleep(5)

        self.log.info('消息队列线程正常退出.')
